import array, time
from machine import Pin
import rp2
from numbers import *


@rp2.asm_pio(sideset_init=rp2.PIO.OUT_LOW, out_shiftdir=rp2.PIO.SHIFT_LEFT, autopull=True, pull_thresh=24)
def ws2812():
    """
    PIO程序，用于控制WS2812 LED。
    通过设置不同的时间参数（T1, T2, T3）来控制LED的亮灭。
    """
    T1 = 2
    T2 = 5
    T3 = 3
    wrap_target()
    label("bitloop")
    out(x, 1)               .side(0)    [T3 - 1]
    jmp(not_x, "do_zero")   .side(1)    [T1 - 1]
    jmp("bitloop")          .side(1)    [T2 - 1]
    label("do_zero")
    nop()                   .side(0)    [T2 - 1]
    wrap()

class WS2812Matrix:
    """
    WS2812 LED Matrix Driver

   为32x8 LED矩阵提供驱动支持。
   处理4个8x8 WS2812模块的串联配置，
   支持单点控制和字符显示。

   注意：
    - 使用RP2040的PIO功能驱动WS2812
    - LED排列为蛇形，需要特殊处理坐标映射
    """
    def __init__(self, pin, width, height):
        """
        初始化WS2812 LED矩阵。
        
        :param pin: 控制LED的GPIO引脚
        :param width: 矩阵的宽度
        :param height: 矩阵的高度
        """
        self.width = width
        self.height = height
        self.panel_width = 8  # 每个8x8面板的宽度
        self.pixels = array.array("I", [0] * (width * height))
        
        # 配置PIO状态机
        self.sm = rp2.StateMachine(0, ws2812, freq=8_000_000, sideset_base=Pin(pin))
        self.sm.active(1)
        
    def _get_pixel_index(self, x, y):
        """
        计算给定坐标在像素数组中的索引。
        
        :param x: 像素的x坐标
        :param y: 像素的y坐标
        :return: 像素在数组中的索引
        """
        # 确定在哪个8x8面板上
        panel_number = x // self.panel_width
        # 计算在当前面板内的x坐标
        local_x = x % self.panel_width
        
        # 计算当前面板内的索引
        if y % 2 == 0:
            # 偶数行：直接从左到右
            local_index = y * self.panel_width + local_x
        else:
            # 奇数行：在每个面板内反转x坐标
            reversed_x = (local_x)
            local_index = y * self.panel_width + reversed_x
        
        # 计算最终索引，注意每个面板的起始位置
        base_index = panel_number * (self.panel_width * self.height)
        return base_index + local_index
    
    def set_pixel(self, x, y, color):
        """
        设置指定坐标的像素颜色。
        
        :param x: 像素的x坐标
        :param y: 像素的y坐标
        :param color: RGB颜色元组
        """
        if 0 <= x < self.width and 0 <= y < self.height:
            rgb = (color[1] << 16) | (color[0] << 8) | color[2]
            self.pixels[self._get_pixel_index(x, y)] = rgb
    
    def get_pixel(self, x, y):
        """
        获取指定坐标的像素颜色。
        
        :param x: 像素的x坐标
        :param y: 像素的y坐标
        :return: RGB颜色元组
        """
        if 0 <= x < self.width and 0 <= y < self.height:
            pixel = self.pixels[self._get_pixel_index(x, y)]
            return ((pixel >> 8) & 0xFF, (pixel >> 16) & 0xFF, pixel & 0xFF)
        return (0, 0, 0)
    
    def show(self):
        """
        将像素数据发送到LED矩阵以更新显示。
        """
        for pixel in self.pixels:
            self.sm.put(pixel, 8)
        time.sleep_ms(10)
        
    def fill(self, color):
        """
        用指定颜色填充整个矩阵。
        
        :param color: RGB颜色元组
        """
        rgb = (color[1] << 16) | (color[0] << 8) | color[2]
        for i in range(len(self.pixels)):
            self.pixels[i] = rgb
            
    def draw_digit(self, digit, x, y, color):
        """
        在指定位置绘制数字。
        
        :param digit: 要绘制的数字
        :param x: 起始x坐标
        :param y: 起始y坐标
        :param color: RGB颜色元组
        """
        pattern = DIGITS[digit]
        for row in range(7):
            for col in range(5):  # 修改为5列
                if pattern[row][col]:
                    self.set_pixel(x + col, y + row, color)
                    
    def draw_colon(self, x, y, color):
        """
        在指定位置绘制冒号。
        
        :param x: 起始x坐标
        :param y: 起始y坐标
        :param color: RGB颜色元组
        """
        self.draw_char(':', x, y, color)
    
    def draw_char(self, char, x, y, color):
        """
        在指定位置绘制字符。
        
        :param char: 要绘制的字符
        :param x: 起始x坐标
        :param y: 起始y坐标
        :param color: RGB颜色元组
        """
        if char in CHARS:
            pattern = CHARS[char]
            for row in range(7):
                for col in range(5):  # 修改为5列
                    if pattern[row][col]:
                        self.set_pixel(x + col, y + row, color) 